'''
Solvers for the linear equations.
Created on 2010-11-17

@author: yun_hua
'''

import numpy as N
import math

class JacobiSolver(object):
    """
    Jacobi method for the linear equations.
    """
    
    def __init__(self, reactor, delta=1e-6, maxiter=100, verbosity=False):
        """
        Constructor
        delta: the tolerance.
        maxiter: the maximum iteration numbers.
        """
        self.grid = reactor.getgrid()
        self.marray = reactor.getmarray()
        self.tarray = reactor.gettarray()
        self.maxiter = maxiter
        self.delta = delta
        self.reactor = reactor
        self.verbosity = verbosity
    
    def run(self):
        """
        Solving the linear equations using the Jacobi method.
        """
        numx = self.grid.shape[1]
        numy = self.grid.shape[0]
        marray = self.marray.copy() # temporary mass (percentage) array
        tarray = self.tarray.copy() # temporary temperature array
        for k in range(self.maxiter):
            for nx in range(1, numx - 1): # we skip the boundaries
                for ny in range(1, numy - 1):
                    maW = self.reactor.getaW(nx, ny, False)
                    maE = self.reactor.getaE(nx, ny, False)
                    maS = self.reactor.getaS(nx, ny, False)
                    maN = self.reactor.getaN(nx, ny, False)
                    maP = self.reactor.getaP(nx, ny, maE, maW, maN, maS, False)
                    msource = self.reactor.getsource(nx, ny, False)
                    marray[ny, nx] = (maW * self.marray[ny, nx - 1] + maE * self.marray[ny, nx + 1] + \
                                    maS * self.marray[ny - 1, nx] + maN * self.marray[ny + 1, nx] + msource) / maP
                    
                    taW = self.reactor.getaW(nx, ny, True)
                    taE = self.reactor.getaE(nx, ny, True)
                    taS = self.reactor.getaS(nx, ny, True)
                    taN = self.reactor.getaN(nx, ny, True)
                    taP = self.reactor.getaP(nx, ny, taE, taW, taN, taS, True)
                    tsource = self.reactor.getsource(nx, ny, True)
                    tarray[ny, nx] = (taW * self.tarray[ny, nx - 1] + taE * self.tarray[ny, nx + 1] + \
                                    taS * self.tarray[ny - 1, nx] + taN * self.tarray[ny + 1, nx] + tsource) / taP
                                    
            merr = abs(self.norm(marray - self.marray))
            terr = abs(self.norm(tarray - self.tarray))
            mrelerr = merr / (self.norm(marray) - N.finfo(float).eps)
            trelerr = terr / (self.norm(tarray) - N.finfo(float).eps)
        
            self.marray = marray.copy()
            self.tarray = tarray.copy()
            self.reactor.setmarray(self.marray)
            self.reactor.settarray(self.tarray)
            if (merr < self.delta or mrelerr < self.delta) and (terr < self.delta or trelerr < self.delta):
                break
            if k % 10 == 0 and self.verbosity:
                print "Iteration number -> %d, merr(%f), terr(%f), mrelerr(%f), trelerr(%f)"  % \
                    (k, merr, terr, mrelerr, trelerr)     
        
    def norm(self, a):        
        """
        Euclidean length of a vector x, where x is unwrapped by an array
        a: the array
        """
        return math.sqrt((a ** 2).sum())
    
    def getsol(self):
        return [self.marray, self.tarray]
         
class GaussSeidelSolver(object):
    """
    Jacobi method for the linear equations.
    """
    
    def __init__(self, reactor, delta=1e-6, maxiter=100, verbosity=False):
        """
        Constructor
        delta: the tolerance.
        maxiter: the maximum iteration numbers.
        """
        self.grid = reactor.getgrid()
        self.marray = reactor.getmarray()
        self.tarray = reactor.gettarray()
        self.maxiter = maxiter
        self.delta = delta
        self.reactor = reactor
        self.verbosity = verbosity
    
    def run(self):
        """
        Solving the linear equations using the Jacobi method.
        """
        numx = self.grid.shape[1]
        numy = self.grid.shape[0]
        marray = self.marray.copy() # temporary mass (percentage) array
        tarray = self.tarray.copy() # temporary temperature array
        merrvar = 0. # temporary error for the mass calculation
        terrvar = 0. # temporary error for the temperature calculation
        for k in range(self.maxiter):
            for nx in range(1, numx - 1): # we skip the boundaries
                for ny in range(1, numy - 1):
                    maW = self.reactor.getaW(nx, ny, False)
                    maE = self.reactor.getaE(nx, ny, False)
                    maS = self.reactor.getaS(nx, ny, False)
                    maN = self.reactor.getaN(nx, ny, False)
                    maP = self.reactor.getaP(nx, ny, maE, maW, maN, maS, False)
                    msource = self.reactor.getsource(nx, ny, False)
                    marray[ny, nx] = (maW * self.marray[ny, nx - 1] + maE * self.marray[ny, nx + 1] + \
                                    maS * self.marray[ny - 1, nx] + maN * self.marray[ny + 1, nx] + msource) / maP
                    # safety check
                    if marray[ny, nx] < 1e-6:
                        print "Iteration number->%d, mass percentage became %f, calculation aborted." % (k, marray[ny, nx]) 
                        return                   
                    # Gauss-Seidel iteration method
                    merrvar += (marray[ny, nx] - self.marray[ny, nx]) ** 2 # Norm calculation step1 first
                    self.reactor.updatemass(nx, ny, marray[ny, nx]) # Update the reactor's mass value
                    self.marray[ny, nx] = marray[ny, nx]            # Update the original mass array, a copy of reactor.marray
                    
                    taW = self.reactor.getaW(nx, ny, True)
                    taE = self.reactor.getaE(nx, ny, True)
                    taS = self.reactor.getaS(nx, ny, True)
                    taN = self.reactor.getaN(nx, ny, True)
                    taP = self.reactor.getaP(nx, ny, taE, taW, taN, taS, True)
                    tsource = self.reactor.getsource(nx, ny, True)
                    tarray[ny, nx] = (taW * self.tarray[ny, nx - 1] + taE * self.tarray[ny, nx + 1] + \
                                    taS * self.tarray[ny - 1, nx] + taN * self.tarray[ny + 1, nx] + tsource) / taP
                                    
                    # Gauss-Seidel iteration method
                    terrvar += (tarray[ny, nx] - self.tarray[ny, nx]) ** 2 # Norm calculation step1 first
                    self.reactor.updatetemp(nx, ny, tarray[ny, nx]) # Update the reactor's heat value 
                    self.tarray[ny, nx] = tarray[ny, nx]            # Update the original heat array, a copy of reactor.tarray
            
            merr = abs(math.sqrt(merrvar)) 
            terr = abs(math.sqrt(terrvar)) 
            merrvar = 0. # Norm calculation step2 for mass
            terrvar = 0. # Norm calculation step2 for heat
            mrelerr = merr / (self.norm(marray) - N.finfo(float).eps)
            trelerr = terr / (self.norm(tarray) - N.finfo(float).eps)
        
            self.marray = marray.copy()
            self.tarray = tarray.copy()
            self.reactor.setmarray(self.marray)
            self.reactor.settarray(self.tarray)
            if (merr < self.delta or mrelerr < self.delta) and (terr < self.delta or trelerr < self.delta):
                break
            if k % 10 == 0 and self.verbosity:
                print "Iteration number -> %d, merr(%f), terr(%f), mrelerr(%f), trelerr(%f)"  % \
                    (k, merr, terr, mrelerr, trelerr)     
        
    def norm(self, a):        
        """
        Euclidean length of a vector x, where x is unwrapped by an array
        a: the array
        """
        return math.sqrt((a ** 2).sum())
    
    def getsol(self):
        #return [self.marray, self.tarray]
        self.reactor.fillboundary()
        return [self.reactor.getmarray(), self.reactor.gettarray(), self.reactor.getconvprofile(), 
                self.reactor.getconcentrationprofile()]       
        
        
    